#include "grid_refinement_node.hpp"
#include <pcl/filters/voxel_grid.h>
#include <pcl/filters/statistical_outlier_removal.h>
#include <sstream>
#include <iomanip>

#ifdef _OPENMP
#include <omp.h>
#endif

GridRefinementNode::GridRefinementNode(ros::NodeHandle& nh) : nh_(nh) {
    // 初始化点云指针
    ground_points_.reset(new pcl::PointCloud<pcl::PointXYZI>);
    nonground_points_.reset(new pcl::PointCloud<pcl::PointXYZI>);

    // 读取参数
    nh_.param("grid_resolution", grid_resolution_, 1.0);
    nh_.param("grid_x_min", grid_x_min_, -50.0);
    nh_.param("grid_x_max", grid_x_max_, 50.0);
    nh_.param("grid_y_min", grid_y_min_, -50.0);
    nh_.param("grid_y_max", grid_y_max_, 50.0);
    
    // 初始化性能监控
    max_history_size_ = 10;
    processing_times_.reserve(max_history_size_);
    
    // 初始化ROS接口
    ground_sub_ = nh_.subscribe("/ground_segmentation/ground", 1, &GridRefinementNode::groundCallback, this);
    nonground_sub_ = nh_.subscribe("/segmented_cloud_pure1", 1, &GridRefinementNode::nongroundCallback, this);
    refined_ground_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("/refined_ground", 1);
    refined_nonground_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("/refined_nonground", 1);
    grid_height_pub_ = nh_.advertise<visualization_msgs::MarkerArray>("/grid_height_visualization", 1);
    
    ROS_INFO("Grid Refinement Node initialized with resolution: %.2f", grid_resolution_);
}

std::string GridRefinementNode::getGridIndex(double x, double y) {
    // 使用与Python版本完全一致的逻辑，但转换为字符串键
    int grid_x = static_cast<int>(std::floor(x / grid_resolution_));
    int grid_y = static_cast<int>(std::floor(y / grid_resolution_));
    
    // 转换为字符串键，格式: "x,y"
    std::ostringstream oss;
    oss << grid_x << "," << grid_y;
    return oss.str();
}

void GridRefinementNode::preallocateGrids() {
    // 预分配网格容器大小，提高性能
    // 基于网格范围估算网格数量
    int grid_count_x = static_cast<int>((grid_x_max_ - grid_x_min_) / grid_resolution_) + 1;
    int grid_count_y = static_cast<int>((grid_y_max_ - grid_y_min_) / grid_resolution_) + 1;
    size_t estimated_grids = static_cast<size_t>(grid_count_x * grid_count_y);
    
    ground_grid_.reserve(estimated_grids);
    grid_avg_heights_.reserve(estimated_grids);
}



       void GridRefinementNode::groundCallback(const sensor_msgs::PointCloud2::ConstPtr& msg) {
           start_time_ = std::chrono::high_resolution_clock::now();
           
           // 直接处理地面点云，不等待
           pcl::fromROSMsg(*msg, *ground_points_);
           ROS_INFO("Received ground points: %zu", ground_points_->size());
           
           processGroundPoints();
           logProcessingTime("Ground points processing");
       }

       void GridRefinementNode::nongroundCallback(const sensor_msgs::PointCloud2::ConstPtr& msg) {
           start_time_ = std::chrono::high_resolution_clock::now();
           
           // 直接处理非地面点云，不等待
           pcl::fromROSMsg(*msg, *nonground_points_);
           ROS_INFO("Received nonground points: %zu", nonground_points_->size());
           
           if (!grid_avg_heights_.empty()) {
               processNongroundPoints();
               logProcessingTime("Nonground points processing");
           } else {
               ROS_WARN("No ground grid data available, skipping nonground processing");
           }
       }

void GridRefinementNode::processGroundPoints() {
    auto assignment_start_time = std::chrono::high_resolution_clock::now();
    
    // 按网格分组地面点
    std::unordered_map<std::string, std::vector<double>> grid_ground_points;
    
    for (const auto& point : ground_points_->points) {
        std::string grid_idx = getGridIndex(point.x, point.y);
        grid_ground_points[grid_idx].push_back(point.z);
    }
    
    auto assignment_end_time = std::chrono::high_resolution_clock::now();
    auto assignment_time = std::chrono::duration_cast<std::chrono::microseconds>(assignment_end_time - assignment_start_time).count();
    ROS_INFO("Grid assignment time: %.3f ms", assignment_time / 1000.0);
    
    // 计算每个网格的地面高度（使用更精确的方法）
    auto height_start_time = std::chrono::high_resolution_clock::now();
    
    for (const auto& grid_pair : grid_ground_points) {
        const std::string& grid_idx = grid_pair.first;
        const std::vector<double>& heights = grid_pair.second;
        
        if (heights.size() > 0) {
            // 始终使用最低点作为网格高度，确保与真实地面匹配
            double min_height = *std::min_element(heights.begin(), heights.end());
            grid_avg_heights_[grid_idx] = min_height;
        }
    }
    
    auto height_end_time = std::chrono::high_resolution_clock::now();
    auto height_time = std::chrono::duration_cast<std::chrono::microseconds>(height_end_time - height_start_time).count();
    ROS_INFO("Height calculation time: %.3f ms", height_time / 1000.0);
    
    ROS_INFO("Processed ground points into %zu grids", grid_avg_heights_.size());
    
    // 输出高度统计信息
    if (!grid_avg_heights_.empty()) {
        std::vector<double> all_heights;
        for (const auto& pair : grid_avg_heights_) {
            all_heights.push_back(pair.second);
        }
        
        std::sort(all_heights.begin(), all_heights.end());
        double min_height = all_heights.front();
        double max_height = all_heights.back();
        double median_height = all_heights[all_heights.size() / 2];
        
        ROS_INFO("Grid height stats: min=%.3f, median=%.3f, max=%.3f", 
                 min_height, median_height, max_height);
        
        // 添加计算高度与真实地面点高度的对比信息
        ROS_INFO("Grid height calculation method: always use minimum height for accurate ground matching");
        ROS_INFO("Visualization now uses real ground point heights and displays at ground level");
    }
    
    // 智能清理：如果点云过大，只清理网格数据，保留原始地面点用于发布
    if (ground_points_->size() > 50000) {
        ROS_INFO("Cleared large ground point cloud to free memory");
        // 不清空原始地面点，只清理网格数据
        // ground_points_->clear();  // 注释掉这行，保留原始地面点
    }
    
    auto end_time = std::chrono::high_resolution_clock::now();
    auto processing_time = std::chrono::duration_cast<std::chrono::microseconds>(end_time - assignment_start_time).count();
    ROS_INFO("Ground points processing time: %.3f ms", processing_time / 1000.0);
}

void GridRefinementNode::processNongroundPoints() {
    auto grouping_start_time = std::chrono::high_resolution_clock::now();
    
    // 按网格分组非地面点
    std::unordered_map<std::string, std::vector<double>> grid_nonground_points;
    
    // 预分配容器大小
    grid_nonground_points.reserve(grid_avg_heights_.size());
    
    for (const auto& point : nonground_points_->points) {
        std::string grid_idx = getGridIndex(point.x, point.y);
        grid_nonground_points[grid_idx].push_back(point.z);
    }
    
    auto grouping_end_time = std::chrono::high_resolution_clock::now();
    auto grouping_time = std::chrono::duration_cast<std::chrono::microseconds>(grouping_end_time - grouping_start_time).count();
    ROS_INFO("Grid grouping time: %.3f ms", grouping_time / 1000.0);
    
    // 重新分类非地面点
    auto classification_start_time = std::chrono::high_resolution_clock::now();
    
    const size_t n_points = nonground_points_->size();
    
    // 预分配点云大小，减少动态分配
    pcl::PointCloud<pcl::PointXYZI>::Ptr refined_ground(new pcl::PointCloud<pcl::PointXYZI>);
    pcl::PointCloud<pcl::PointXYZI>::Ptr refined_nonground(new pcl::PointCloud<pcl::PointXYZI>);
    
    // 更激进的内存预分配
    refined_ground->points.reserve(n_points);
    refined_nonground->points.reserve(n_points);
    
    // 批量处理：先收集所有点的分类结果，再一次性添加
    std::vector<std::pair<size_t, bool>> classification_results; // <point_index, is_ground>
    classification_results.reserve(n_points);
    
    // 第一阶段：快速分类，只记录结果，不创建点云
    #ifdef _OPENMP
    #pragma omp parallel for
    #endif
    for (size_t i = 0; i < n_points; ++i) {
        const auto& point = nonground_points_->points[i];
        std::string grid_idx = getGridIndex(point.x, point.y);
        
        bool is_ground = false;
        
        // 检查该网格是否有地面高度数据
        auto ground_height_it = grid_avg_heights_.find(grid_idx);
        if (ground_height_it != grid_avg_heights_.end()) {
            // 网格有地面数据：使用更精确的自适应阈值
            double grid_height = ground_height_it->second;
            
            // 根据网格内点的数量调整阈值
            auto nonground_it = grid_nonground_points.find(grid_idx);
            if (nonground_it != grid_nonground_points.end() && nonground_it->second.size() >= 5) {
                // 如果该网格有足够的非地面点，使用统计方法计算阈值
                std::vector<double> sorted_heights = nonground_it->second;
                std::sort(sorted_heights.begin(), sorted_heights.end());
                
                // 使用非地面点的最低点+0.15作为阈值（更精确）
                double min_nonground_height = sorted_heights[0];
                double adaptive_threshold = std::min(grid_height + 0.15, min_nonground_height + 0.15);
                is_ground = (point.z <= adaptive_threshold);
            } else {
                // 使用保守的阈值：grid_height + 0.15
                is_ground = (point.z <= grid_height + 0.15);
            }
        } else {
            // 网格没有地面数据：使用该网格内非地面点的最低点+0.15作为过滤阈值
            auto nonground_it = grid_nonground_points.find(grid_idx);
            if (nonground_it != grid_nonground_points.end() && !nonground_it->second.empty()) {
                double grid_min_height = *std::min_element(nonground_it->second.begin(), 
                                                        nonground_it->second.end());
                double grid_threshold = grid_min_height + 0.15;
                is_ground = (point.z <= grid_threshold);
            } else {
                // 如果该网格内没有非地面点，保持为非地面点（与Python版本一致）
                is_ground = false;
            }
        }
        
        classification_results.emplace_back(i, is_ground);
    }
    
    // 第二阶段：批量创建点云，减少内存分配次数
    // 先统计地面点数量
    size_t ground_count = 0;
    for (const auto& result : classification_results) {
        if (result.second) ground_count++;
    }
    
    // 调整预分配大小 - 包含原始地面点
    refined_ground->points.resize(ground_count + ground_points_->size());
    refined_nonground->points.resize(n_points - ground_count);
    
    // 批量填充点云
    size_t ground_idx = 0;
    size_t nonground_idx = 0;
    
    // 首先添加原始地面点（强度设为50）
    for (const auto& point : ground_points_->points) {
        auto& new_point = refined_ground->points[ground_idx++];
        new_point.x = point.x;
        new_point.y = point.y;
        new_point.z = point.z;
        new_point.intensity = 50.0;
    }
    
    ROS_INFO("Added %zu original ground points", ground_points_->size());
    ROS_INFO("After adding original ground points, ground_idx = %zu", ground_idx);
    
    // 然后添加重新分类的点
    for (const auto& result : classification_results) {
        const auto& point = nonground_points_->points[result.first];
        
        if (result.second) {
            // 重新分类为地面点（强度设为50）
            auto& new_point = refined_ground->points[ground_idx++];
            new_point.x = point.x;
            new_point.y = point.y;
            new_point.z = point.z;
            new_point.intensity = 50.0;
        } else {
            // 保持为非地面点（强度设为100）
            auto& new_point = refined_nonground->points[nonground_idx++];
            new_point.x = point.x;
            new_point.y = point.y;
            new_point.z = point.z;
            new_point.intensity = 100.0;
        }
    }
    
    auto classification_end_time = std::chrono::high_resolution_clock::now();
    auto classification_time = std::chrono::duration_cast<std::chrono::microseconds>(classification_end_time - classification_start_time).count();
    ROS_INFO("Point classification time: %.3f ms", classification_time / 1000.0);
    
    // 发布重新分类后的点云
    publishPointCloud(refined_ground, refined_ground_pub_, "base_laser");
    publishPointCloud(refined_nonground, refined_nonground_pub_, "base_laser");
    
    // 发布网格高度可视化
    publishGridHeightVisualization();
    
    ROS_INFO("Refined classification: %zu ground points, %zu nonground points", 
             refined_ground->size(), refined_nonground->size());
}

void GridRefinementNode::publishGridHeightVisualization() {
    if (grid_avg_heights_.empty()) {
        return;
    }
    
    visualization_msgs::MarkerArray marker_array;
    int marker_id = 0;
    
    // 统计真实地面点高度
    std::unordered_map<std::string, std::vector<double>> grid_real_heights;
    
    // 从原始地面点计算每个网格的真实高度
    for (const auto& point : ground_points_->points) {
        std::string grid_idx = getGridIndex(point.x, point.y);
        grid_real_heights[grid_idx].push_back(point.z);
    }
    
    // 只遍历有地面数据的网格
    for (const auto& grid_pair : grid_avg_heights_) {
        const std::string& grid_key = grid_pair.first;
        double calculated_height = grid_pair.second;
        
        // 获取该网格的真实地面点高度
        double real_height = calculated_height; // 默认使用计算高度
        auto real_height_it = grid_real_heights.find(grid_key);
        if (real_height_it != grid_real_heights.end() && !real_height_it->second.empty()) {
            // 使用真实地面点的中位数高度
            std::vector<double> heights = real_height_it->second;
            std::sort(heights.begin(), heights.end());
            
            if (heights.size() % 2 == 0) {
                real_height = (heights[heights.size()/2 - 1] + heights[heights.size()/2]) / 2.0;
            } else {
                real_height = heights[heights.size()/2];
            }
        }
        
        // 解析字符串键，获取x和y
        std::istringstream iss(grid_key);
        int grid_x, grid_y;
        char comma;
        iss >> grid_x >> comma >> grid_y;
        
        // 计算网格边界坐标
        double x_min = grid_x * grid_resolution_;
        double x_max = (grid_x + 1) * grid_resolution_;
        double y_min = grid_y * grid_resolution_;
        double y_max = (grid_y + 1) * grid_resolution_;
        
        // 创建4条边的Marker（空心网格），使用真实高度
        std::vector<std::pair<Eigen::Vector3d, Eigen::Vector3d>> edges = {
            {{(x_min + x_max)/2, y_min, real_height}, {x_max - x_min, 0.05, 0.1}},  // 下边
            {{(x_min + x_max)/2, y_max, real_height}, {x_max - x_min, 0.05, 0.1}},  // 上边
            {{x_min, (y_min + y_max)/2, real_height}, {0.05, y_max - y_min, 0.1}},  // 左边
            {{x_max, (y_min + y_max)/2, real_height}, {0.05, y_max - y_min, 0.1}}   // 右边
        };
        
        for (size_t i = 0; i < edges.size(); ++i) {
            visualization_msgs::Marker marker;
            marker.header.frame_id = "base_laser";
            marker.header.stamp = ros::Time::now();
            marker.ns = "grid_height";
            marker.id = marker_id * 4 + i;
            marker.type = visualization_msgs::Marker::CUBE;
            marker.action = visualization_msgs::Marker::ADD;
            
            // 设置边缘位置和尺寸，使用真实高度
            marker.pose.position.x = edges[i].first.x();
            marker.pose.position.y = edges[i].first.y();
            marker.pose.position.z = edges[i].first.z();
            marker.pose.orientation.w = 1.0;
            
            // 设置边缘尺寸
            marker.scale.x = edges[i].second.x();
            marker.scale.y = edges[i].second.y();
            marker.scale.z = edges[i].second.z();
            
            // 根据真实高度设置颜色（高度越高，颜色越红）
            double height_normalized = std::max(0.0, std::min(1.0, (real_height + 2.0) / 4.0));
            marker.color.r = height_normalized;
            marker.color.g = 1.0 - height_normalized;
            marker.color.b = 0.0;
            marker.color.a = 0.8;
            
            marker_array.markers.push_back(marker);
        }
        
        marker_id++;
    }
    
    // 发布MarkerArray
    grid_height_pub_.publish(marker_array);
    
    // 输出高度匹配信息
    if (marker_id > 0) {
        ROS_DEBUG("Published %d grid height markers using real ground point heights", marker_id);
    }
}

void GridRefinementNode::publishPointCloud(const pcl::PointCloud<pcl::PointXYZI>::Ptr& cloud, 
                                          ros::Publisher& publisher, const std::string& frame_id) {
    if (cloud->empty()) {
        return;
    }
    
    sensor_msgs::PointCloud2 cloud_msg;
    pcl::toROSMsg(*cloud, cloud_msg);
    cloud_msg.header.frame_id = frame_id;
    cloud_msg.header.stamp = ros::Time::now();
    
    publisher.publish(cloud_msg);
}

void GridRefinementNode::logProcessingTime(const std::string& operation_name) {
    auto end_time = std::chrono::high_resolution_clock::now();
    auto processing_time = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time_).count();
    double time_ms = processing_time / 1000.0;
    
    ROS_INFO("%s time: %.3f ms", operation_name.c_str(), time_ms);
    
    // 监控性能
    monitorPerformance(time_ms);
}

void GridRefinementNode::monitorPerformance(double processing_time) {
    processing_times_.push_back(processing_time);
    
    // 保持历史记录大小
    if (processing_times_.size() > max_history_size_) {
        processing_times_.erase(processing_times_.begin());
    }
    

}



int main(int argc, char** argv) {
    ros::init(argc, argv, "grid_refinement_node");
    ros::NodeHandle nh("~");
    
    try {
        GridRefinementNode node(nh);
        ROS_INFO("Grid Refinement Node is running...");
        ros::spin();
    } catch (const std::exception& e) {
        ROS_ERROR("Exception in main: %s", e.what());
        return 1;
    }
    
    return 0;
} 